import matplotlib.pyplot as plt
import numpy as np
from skimage import data, feature, exposure
from skimage.color import rgb2gray
from matplotlib.patches import Rectangle

# 加载示例图像
image = data.astronaut()

# 将图像转换为灰度图
gray_image = rgb2gray(image)

# HOG参数设置
orientations = 8
pixels_per_cell = (32, 32)
cells_per_block = (1, 1)

# 计算HOG特征
fd, hog_image = feature.hog(gray_image, orientations=orientations,
                            pixels_per_cell=pixels_per_cell,
                            cells_per_block=cells_per_block,
                            visualize=True)

# 获取图像尺寸和cell数量
rows, cols = gray_image.shape[:2]
cells_y, cells_x = (rows // pixels_per_cell[0], cols // pixels_per_cell[1])

# 重塑特征向量以获取每个cell的HOG特征
fd_reshaped = fd.reshape(cells_y, cells_x, cells_per_block[0], cells_per_block[1], orientations)
# 取每个cell的第一个block的HOG特征
cell_hog = fd_reshaped[:, :, 0, 0, :]

# 创建一个3列的图像显示
fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(18, 6))

# 显示原始图像
ax1.imshow(image)
ax1.set_title('原始图像')
ax1.axis('off')

# 显示标记了cell的中间图像（使用灰度图）
ax2.imshow(gray_image, cmap=plt.cm.gray)
ax2.set_title('标记了Cell的灰度图像')

# 绘制cell网格
for i in range(cells_y):
    for j in range(cells_x):
        # 绘制cell边界
        rect = Rectangle((j * pixels_per_cell[1], i * pixels_per_cell[0]),
                         pixels_per_cell[1], pixels_per_cell[0],
                         linewidth=1, edgecolor='r', facecolor='none')
        ax2.add_patch(rect)

ax2.axis('off')

# 创建空白背景图像
max_hog_image = np.zeros((rows, cols))

# 绘制每个cell中幅值最大的梯度线条
for i in range(cells_y):
    for j in range(cells_x):
        # 找到最大幅值的梯度方向
        max_orientation_index = np.argmax(cell_hog[i, j])
        max_magnitude = cell_hog[i, j][max_orientation_index]

        # 计算线条的角度
        angle = (max_orientation_index * 180 / orientations) + 90 / orientations
        angle_rad = np.deg2rad(angle)

        # 计算cell的中心位置
        center_x = int(j * pixels_per_cell[1] + pixels_per_cell[1] / 2)
        center_y = int(i * pixels_per_cell[0] + pixels_per_cell[0] / 2)

        # 计算线条的长度
        line_length = int(max_magnitude * pixels_per_cell[0] * 0.5)

        # 计算线条的端点
        dx = int(line_length * np.cos(angle_rad))
        dy = int(line_length * np.sin(angle_rad))

        # 绘制线条到空白图像上，用亮度表示幅值
        for t in np.linspace(-1, 1, 20):
            x = int(center_x + t * dx)
            y = int(center_y + t * dy)
            if 0 <= x < cols and 0 <= y < rows:
                max_hog_image[y, x] = max_magnitude

# 归一化图像
max_hog_image = max_hog_image / np.max(max_hog_image)

# 显示只包含最大幅值梯度线条的图像
ax3.imshow(max_hog_image, cmap=plt.cm.gray)
ax3.set_title('只显示最大幅值梯度线条的图像')
ax3.axis('off')

plt.tight_layout()
plt.show()
